from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from typing import List, Optional
import nltk

# 下载必要的NLTK资源
import ssl

try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context

# 检查并下载punkt分词器
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')

class DocumentSplitter:
    """文档文本分割器，使用递归字符分割策略保持语义完整性"""
    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        separators: Optional[List[str]] = None,
        add_start_index: bool = True
    ):
        """
        初始化文本分割器
        :param chunk_size: 文本块大小（字符数）
        :param chunk_overlap: 文本块重叠部分大小
        :param separators: 自定义分隔符列表
        :param add_start_index: 是否添加原始文档中的起始索引
        """
        # 默认分隔符，优先使用段落分隔，再句子，再单词
        default_separators = [
            "\n\n", "\n", "\r\n\r\n", "\r\n",  # 段落分隔
            ". ", "! ", "? ", "; ", ": ",  # 句子分隔
            " ", ""
        ]

        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators or default_separators,
            add_start_index=add_start_index,
            keep_separator=False,
            length_function=len
        )

    def split_document(self, document: Document) -> List[Document]:
        """
        分割单个文档
        :param document: 要分割的文档对象
        :return: 分割后的文档块列表
        """
        if not document or not document.page_content:
            return []
        return self.splitter.split_documents([document])

    def split_documents(self, documents: List[Document]) -> List[Document]:
        """
        分割多个文档
        :param documents: 文档对象列表
        :return: 所有分割后的文档块列表
        """
        if not documents:
            return []
        return self.splitter.split_documents(documents)

    def create_split_chunks(self, text: str, metadata: Optional[dict] = None) -> List[Document]:
        """
        直接从文本创建分割块
        :param text: 原始文本
        :param metadata: 文本元数据
        :return: 分割后的文档块列表
        """
        if not text:
            return []
        document = Document(page_content=text, metadata=metadata or {})
        return self.split_document(document)

    def adjust_chunk_parameters(self, **kwargs) -> None:
        """动态调整分割器参数"""
        for key, value in kwargs.items():
            if hasattr(self.splitter, key):
                setattr(self.splitter, key, value)

if __name__ == "__main__":
    # 示例用法
    splitter = DocumentSplitter(
        chunk_size=500,
        chunk_overlap=100
    )

    # 示例文档
    sample_document = Document(
        page_content="""LangChain是一个强大的框架，用于构建基于语言模型的应用程序。\n\n"
        "它提供了一系列工具和接口，可以轻松地将语言模型与外部数据源集成，\n"
        "实现检索增强生成（RAG）等高级功能。\n\n"
        "RAG技术允许模型在生成回答时引用外部知识库，提高回答的准确性和可靠性。\n"
        "在本项目中，我们将使用LangChain构建一个完整的RAG系统，包括文档加载、文本分割、\n"
        "向量存储、检索和生成等核心模块。\n\n"
        "文本分割是RAG流程中的关键步骤，它将长文档分解为适合向量化的小块，\n"
        "同时尽量保持文本的语义完整性。这对于后续的检索准确性至关重要。"
        """,
        metadata={"source": "example.txt", "page": 1}
    )

    # 分割单个文档
    chunks = splitter.split_document(sample_document)
    print(f"原始文档长度: {len(sample_document.page_content)}字符")
    print(f"分割后块数: {len(chunks)}")
    for i, chunk in enumerate(chunks):
        print(f"\n块 {i+1}: {len(chunk.page_content)}字符")
        print(f"内容: {chunk.page_content[:100]}...")
        print(f"元数据: {chunk.metadata}")